from sqlalchemy import create_engine, MetaData, Table, select, and_, update, insert, desc, asc
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.engine import Engine
from typing import List, Dict, Any, Optional

import cfg


class DatabaseHandler:
    def __init__(self, connection_string: str):
        """
        初始化 DatabaseHandler。

        :param connection_string: 数据库连接字符串，例如：
            'mysql+pymysql://username:password@localhost:3306/database_name'
        """
        self.engine: Engine = create_engine(connection_string, echo=False, future=True)
        self.metadata = MetaData()
        self.metadata.reflect(bind=self.engine)  # 反射数据库中所有表
        self._table_cache = self.metadata.tables  # 缓存表对象

    def _get_table(self, table_name: str) -> Table:
        """
        根据表名获取 SQLAlchemy Table 对象。

        :param table_name: 表名
        :return: Table 对象
        :raises ValueError: 如果表不存在
        """
        if table_name not in self._table_cache:
            try:
                # 尝试反射特定表
                table = Table(table_name, self.metadata, autoload_with=self.engine)
                self._table_cache[table_name] = table
            except SQLAlchemyError:
                raise ValueError(f"表 '{table_name}' 不存在。")
        return self._table_cache[table_name]

    def insert(self, table_name: str, params: Dict[str, Any]) -> None:
        """
        安全地插入一条记录到指定的表中。

        :param table_name: 目标表名
        :param params: 一个字典，键为列名，值为插入的值
        """
        table = self._get_table(table_name)

        # 验证列名
        invalid_columns = set(params.keys()) - set(table.c.keys())
        if invalid_columns:
            raise ValueError(f"无效的列名在 params 中: {', '.join(invalid_columns)}")

        stmt = insert(table).values(**params)

        try:
            with self.engine.connect() as conn:
                result = conn.execute(stmt)
                conn.commit()
                lastrowid = result.lastrowid
                print(f"新增的 ID 是: {lastrowid}")
                return lastrowid
        except SQLAlchemyError as e:
            print(f"插入操作失败: {e}")
            raise

    def update(self, table_name: str, params: Dict[str, Any], where_params: Dict[str, Any]) -> None:
        """
        安全地更新指定表中的记录。

        :param table_name: 目标表名
        :param params: 一个字典，键为列名，值为更新后的值
        :param where_params: 一个字典，键为列名，值为筛选条件的值
        """
        table = self._get_table(table_name)

        # 验证列名
        all_keys = set(params.keys()).union(set(where_params.keys()))
        invalid_columns = all_keys - set(table.c.keys())
        if invalid_columns:
            raise ValueError(f"无效的列名: {', '.join(invalid_columns)}")

        # 构建 WHERE 条件
        conditions = [table.c[key] == value for key, value in where_params.items()]
        stmt = update(table).where(and_(*conditions)).values(**params)

        try:
            with self.engine.connect() as conn:
                result = conn.execute(stmt)
                conn.commit()
                if result.rowcount == 0:
                    print("未更新任何记录。")
        except SQLAlchemyError as e:
            print(f"更新操作失败: {e}")
            raise

    def query(
            self,
            table_name: str,
            columns: Optional[List[str]] = None,
            where_params: Optional[Dict[str, Any]] = None,
            order_by: Optional[List[str]] = None,
            limit: Optional[int] = None,
            offset: Optional[int] = None
    ) -> List[Dict[str, Any]]:
        """
        安全地查询指定表中的记录，支持更多的筛选条件。

        :param table_name: 目标表名
        :param columns: 一个列表，指定要查询的列名。如果为 None 或空列表，则查询所有列。
        :param where_params: 一个字典，键为列名，值为筛选条件的值。支持操作符，如 'column__like': '%abc%'
        :param order_by: 一个列表，指定排序的列和方向，例如 ["name ASC", "age DESC"]
        :param limit: 一个整数，指定返回的最大记录数
        :param offset: 一个整数，指定记录的偏移量
        :return: 查询结果的列表，每个记录为字典
        """
        table = self._get_table(table_name)

        # 选择列
        if not columns:
            stmt = select(table)
        else:
            invalid_columns = set(columns) - set(table.c.keys())
            if invalid_columns:
                raise ValueError(f"无效的列名在 columns 中: {', '.join(invalid_columns)}")
            selected_columns = [table.c[col] for col in columns]
            stmt = select(*selected_columns)

        # 添加 WHERE 条件
        if where_params:
            conditions = []
            for key, value in where_params.items():
                if '__' in key:
                    col_name, op = key.split('__', 1)
                    if col_name not in table.c:
                        raise ValueError(f"无效的列名在 where_params 中: '{col_name}'")
                    column = table.c[col_name]
                    if op == 'like':
                        conditions.append(column.like(value))
                    elif op == 'ilike':
                        conditions.append(column.ilike(value))
                    elif op == 'gt':
                        conditions.append(column > value)
                    elif op == 'lt':
                        conditions.append(column < value)
                    elif op == 'gte':
                        conditions.append(column >= value)
                    elif op == 'lte':
                        conditions.append(column <= value)
                    elif op == 'ne':
                        conditions.append(column != value)
                    elif op == 'in':
                        if not isinstance(value, list):
                            raise ValueError(f"操作符 'in' 的值必须是列表。")
                        conditions.append(column.in_(value))
                    elif op == 'notin':
                        if not isinstance(value, list):
                            raise ValueError(f"操作符 'notin' 的值必须是列表。")
                        conditions.append(~column.in_(value))
                    else:
                        raise ValueError(f"不支持的操作符: '{op}'")
                else:
                    if key not in table.c:
                        raise ValueError(f"无效的列名在 where_params 中: '{key}'")
                    conditions.append(table.c[key] == value)
            stmt = stmt.where(and_(*conditions))

        # 添加 ORDER BY
        if order_by:
            order_clauses = []
            for item in order_by:
                parts = item.strip().split()
                if len(parts) == 1:
                    col_name, direction = parts[0], "asc"
                elif len(parts) == 2:
                    col_name, direction = parts
                    direction = direction.lower()
                else:
                    raise ValueError(f"无效的 order_by 格式: '{item}'")

                if col_name not in table.c:
                    raise ValueError(f"无效的排序列: '{col_name}'")

                column = table.c[col_name]
                if direction == "asc":
                    order_clauses.append(asc(column))
                elif direction == "desc":
                    order_clauses.append(desc(column))
                else:
                    raise ValueError(f"无效的排序方向: '{direction}'，应为 'ASC' 或 'DESC'。")

            stmt = stmt.order_by(*order_clauses)

        # 添加 LIMIT 和 OFFSET
        if limit is not None:
            if not isinstance(limit, int) or limit < 0:
                raise ValueError("limit 必须是非负整数")
            stmt = stmt.limit(limit)
            if offset is not None:
                if not isinstance(offset, int) or offset < 0:
                    raise ValueError("offset 必须是非负整数")
                stmt = stmt.offset(offset)

        try:
            with self.engine.connect() as conn:
                result = conn.execute(stmt)
                rows = result.mappings().all()
                return [dict(row) for row in rows]
        except SQLAlchemyError as e:
            print(f"查询操作失败: {e}")
            raise

    def __del__(self):
        """
        析构方法，确保释放数据库连接资源。
        """
        self.engine.dispose()


if __name__ == '__main__':
    db_url = f'mysql+pymysql://{cfg.user}:{cfg.password}@{cfg.host}:{cfg.port}/bank'
    handler = DatabaseHandler(db_url)
    # result = handler.query('ntqmdlstz', ['id', 'busmod', 'modals'])
    # print(result)
    handler.insert('ntqmdlstz', {'busmod': '1', 'modals': '1'})
